package com.chengxiaoxiao.ai_demo.service.impl;

import com.chengxiaoxiao.ai_demo.config.FileProperties;
import com.chengxiaoxiao.ai_demo.config.RagProperties;
import com.chengxiaoxiao.ai_demo.entity.KnowledgeFile;
import com.chengxiaoxiao.ai_demo.service.DocumentEmbeddingService;
import com.chengxiaoxiao.ai_demo.vo.vo.KnowledgeSearchVo;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.parser.apache.tika.ApacheTikaDocumentParser;
import dev.langchain4j.data.document.splitter.DocumentSplitters;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.Tokenizer;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.store.embedding.EmbeddingMatch;
import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.filter.Filter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Async;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.List;

import static com.chengxiaoxiao.ai_demo.config.AiAutoConfiguration.META_COLUMN_KNOWLEDGE_BASE_ID;
import static dev.langchain4j.store.embedding.filter.MetadataFilterBuilder.metadataKey;


/**
 * 文件嵌入 服务实现类
 *
 * @Description:
 * @Author: Cheng Xiaoxiao
 * @Date: 2025-03-17 15:06:21
 */
@Slf4j
@Service
@RequiredArgsConstructor
public class DocumentEmbeddingServiceImpl implements DocumentEmbeddingService {
    /**
     * 元数据列：知识库文件ID
     */
    public static final String META_COLUMN_KNOWLEDGE_FILE_ID = "knowledge_file_id";

    final FileProperties fileProperties;
    final EmbeddingModel embeddingModel;
    final EmbeddingStore embeddingStore;
    final Tokenizer tokenizer;
    final RagProperties ragProperties;

    /**
     * 嵌入知识库文件
     *
     * @param knowledgeBaseId 知识库ID
     * @param knowledgeFiles  知识库文件
     */
    @Async
    @Override
    public void embeddingDocument(String knowledgeBaseId, List<KnowledgeFile> knowledgeFiles) {

        for (KnowledgeFile knowledgeFile : knowledgeFiles) {
            String fileName = fileProperties.getSavePath() + knowledgeFile.getStorePath();
            Document doc = FileSystemDocumentLoader.loadDocument(fileName, new ApacheTikaDocumentParser());

            EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                    .documentSplitter(DocumentSplitters.recursive(300, 30, tokenizer))
                    .documentTransformer(document -> {
                        document.metadata().remove("absolute_directory_path");
                        document.metadata().put("file_name", knowledgeFile.getFileName());
                        document.metadata().put(META_COLUMN_KNOWLEDGE_BASE_ID, knowledgeBaseId);
                        document.metadata().put(META_COLUMN_KNOWLEDGE_FILE_ID, knowledgeFile.getId());
                        return document;
                    })
                    .embeddingStore(embeddingStore)
                    .embeddingModel(embeddingModel)
                    .build();
            ingestor.ingest(doc);
        }

        log.info("文档嵌入完毕");
    }

    /**
     * 根据知识库ID和关键词搜索
     *
     * @param id  知识库ID
     * @param key 关键词
     * @return 搜索结果
     */
    @Override
    public List<KnowledgeSearchVo> search(String id, String key) {

        Filter filter = metadataKey(META_COLUMN_KNOWLEDGE_BASE_ID).isEqualTo(id);

        EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
                .filter(filter)
                .queryEmbedding(embeddingModel.embed(key).content())
                .maxResults(ragProperties.getMaxResults())
                .minScore(ragProperties.getMinScore())
                .build();
        List<EmbeddingMatch<TextSegment>> relevant = embeddingStore.search(request).matches();

        List<KnowledgeSearchVo> result = new ArrayList<>();
        for (EmbeddingMatch<TextSegment> textSegmentEmbeddingMatch : relevant) {
            Double score = textSegmentEmbeddingMatch.score();
            String text = textSegmentEmbeddingMatch.embedded().text();
            result.add(new KnowledgeSearchVo().setText(text).setScore(score));
        }

        return result;
    }

    /**
     * 根据知识库ID删除
     *
     * @param knowledgeBaseId 知识库ID
     */
    @Override
    public void deleteByKnowledgeBaseId(String knowledgeBaseId) {
        Filter filter = metadataKey(META_COLUMN_KNOWLEDGE_BASE_ID).isEqualTo(knowledgeBaseId);
        embeddingStore.removeAll(filter);
    }
}